package mywebsocket

import (
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

type Connection struct {
	conn       *websocket.Conn
	lastActive int64
	mu         sync.Mutex
}

type Manager struct {
	Connections sync.Map
}

func NewManager() *Manager {
	return &Manager{}
}

// Gin中间件：将Manager实例注入上下文
func (m *Manager) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set("wsManager", m)
		c.Next()
	}
}

func (m *Manager) Add(conn *websocket.Conn) {
	m.Connections.Store(conn, &Connection{
		conn:       conn,
		lastActive: time.Now().Unix(),
	})
}

func (m *Manager) Remove(conn *websocket.Conn) {
	m.Connections.Delete(conn)
}

func (m *Manager) GetAll() []*websocket.Conn {
	var conns []*websocket.Conn
	m.Connections.Range(func(k, _ interface{}) bool {
		conns = append(conns, k.(*websocket.Conn))
		return true
	})
	return conns
}

func (m *Manager) UpdateActive(conn *websocket.Conn) {
	if v, ok := m.Connections.Load(conn); ok {
		conn := v.(*Connection)
		conn.mu.Lock()
		defer conn.mu.Unlock()
		conn.lastActive = time.Now().Unix()
	}
}
